package site.https;

import java.net.*;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class DNSResolver {
    public static void main(String[] args) {
        String[] domains = {"www.google.com", "facebook.com"};
        String dnsServer = "8.8.8.8"; // Google Public DNS
        int port = 53; // DNS server port
        
        try (DatagramSocket socket = new DatagramSocket()) {
            InetAddress serverAddr = InetAddress.getByName(dnsServer);
            socket.connect(new InetSocketAddress(serverAddr, port));
            
            for (String domain : domains) {
                byte[] query = createDNSQuery(domain);
                DatagramPacket queryPacket = new DatagramPacket(query, query.length, serverAddr, port);
                socket.send(queryPacket);
                
                byte[] response = new byte[1024];
                DatagramPacket responsePacket = new DatagramPacket(response, response.length);
                socket.receive(responsePacket);
                
                List<String> addresses = parseDNSResponse(responsePacket.getData());
                System.out.println("IP addresses for " + domain + ":");
                for (String address : addresses) {
                    System.out.println(address);
                }
            }
        } catch (Exception ex) {
            System.out.println("Error: " + ex.getMessage());
        }
    }

    private static byte[] createDNSQuery(String domain) {
        ByteBuffer buffer = ByteBuffer.allocate(512);

        // DNS header
        buffer.putShort((short) 0x1234); // DNS transaction ID
        buffer.putShort((short) 0x0000); // DNS flags (standard query)
        buffer.putShort((short) 0x0001); // DNS question count
        buffer.putShort((short) 0x0000); // DNS answer count
        buffer.putShort((short) 0x0000); // DNS authority count
        buffer.putShort((short) 0x0000); // DNS additional count

        // DNS question section
        String[] labels = domain.split("\\.");
        for (String label : labels) {
            byte[] labelBytes = label.getBytes();
            buffer.put((byte) labelBytes.length);
            buffer.put(labelBytes);
        }
        buffer.put((byte) 0x00); // end of domain name
        buffer.putShort((short) 0x0001); // DNS query type (A record)
        buffer.putShort((short) 0x0001); // DNS query class (IN)

        return Arrays.copyOfRange(buffer.array(), 0, buffer.position());
    }


    public static List<String> parseDNSResponse(byte[] response) throws UnknownHostException {
        List<String> ips = new ArrayList<>();
        if (response.length < 12) {
            throw new UnknownHostException("Invalid DNS response");
        }
        int questionCount = (((int) response[4]) << 8) + ((int) response[5]);
        int answerCount = (((int) response[6]) << 8) + ((int) response[7]);
        if (questionCount != 1) {
            throw new UnknownHostException("Invalid DNS response: question count is not 1");
        }
        if ((response[2] & 0x80) == 0) {
            throw new UnknownHostException("DNS response is not a response");
        }
        if ((response[3] & 0x0f) != 0) {
            throw new UnknownHostException("Error in DNS response: " + (response[3] & 0x0f));
        }
        int pos = 12;
        for (int i = 0; i < answerCount; i++) {
            int type = (((int) response[pos + 2]) << 8) + ((int) response[pos + 3]);
            int dataLength = (((int) response[pos + 8]) << 8) + ((int) response[pos + 9]);
            if ( dataLength == 4) {
                StringBuilder sb = new StringBuilder();
                for (int j = 0; j < 4; j++) {
                    if (j > 0) {
                        sb.append(".");
                    }
                    sb.append(response[pos + 10 + j] & 0xff);
                }
                ips.add(sb.toString());
            }
            pos += 10 + dataLength;
        }
        if (ips.isEmpty()) {
            throw new UnknownHostException("No A records found in DNS response");
        }
        return ips;
    }


}
